package com.babel.util.ciphers;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.Charset;
import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;

import org.apache.commons.codec.binary.Base64;

/**
 * @author Kun
 * @date 2017-6-2 16:52:42
 */
public class RSA {

	private static final String ALGORITHM_RSA = "RSA";

	private static final String DEFAULT_ENCODE = "UTF-8";

	private static final int KEY_SIZE = 2048;

	/**
	 * 加密缓冲区
	 */
	private static final int ENCRYPT_BUFFER_SIZE = 245;

	/**
	 * 解密缓冲区
	 */
	private static final int DECRYPT_BUFFER_SIZE = 256;

	/**
	 * 解密
	 * 
	 * @param privateKey
	 * @param encryptedMessage
	 * @return
	 */
	public static String decrypt(String privateKey, String encryptedMessage) {
		return new String(decrypt(getPrivateKey(privateKey), decodeBase64(encryptedMessage)),
				Charset.forName(DEFAULT_ENCODE));
	}

	/**
	 * 解密
	 * 
	 * @param privateKey
	 * @param encryptedMessage
	 * @return
	 */
	private static byte[] decrypt(PrivateKey privateKey, byte[] encryptedMessage) {
		try (InputStream in = new ByteArrayInputStream(encryptedMessage);
				ByteArrayOutputStream out = new ByteArrayOutputStream();) {
			Cipher cipher = Cipher.getInstance(ALGORITHM_RSA);
			cipher.init(Cipher.DECRYPT_MODE, privateKey);
			byte[] buf = new byte[DECRYPT_BUFFER_SIZE];
			int readCount = 0;
			while ((readCount = in.read(buf)) != -1) {
				byte[] block = null;
				if (buf.length == readCount) {
					block = buf;
				} else {
					block = new byte[readCount];
					System.arraycopy(buf, 0, block, 0, readCount);
				}
				out.write(cipher.doFinal(block));
				out.flush();
			}
			return out.toByteArray();
		} catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException | IllegalBlockSizeException
				| BadPaddingException | IOException e) {
			throw new IllegalStateException(e);
		}
	}

	/**
	 * 加密
	 * 
	 * @param publicKey
	 * @param plainMessage
	 * @return
	 */
	public static String encrypt(String publicKey, String plainMessage) {
		return base64(encrypt(getPublicKey(publicKey), plainMessage.getBytes()));
	}

	private static PrivateKey getPrivateKey(String privateKey) {
		try {
			PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(decodeBase64(privateKey));
			KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM_RSA);
			return keyFactory.generatePrivate(keySpec);
		} catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
			throw new IllegalStateException(e);
		}
	}

	private static PublicKey getPublicKey(String publicKey) {
		try {
			X509EncodedKeySpec keySpec = new X509EncodedKeySpec(decodeBase64(publicKey));
			KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM_RSA);
			return keyFactory.generatePublic(keySpec);
		} catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
			throw new IllegalStateException(e);
		}
	}

	/**
	 * 加密
	 * 
	 * @param publicKey
	 * @param plainMessage
	 * @return
	 */
	private static byte[] encrypt(PublicKey publicKey, byte[] plainMessage) {
		try (InputStream in = new ByteArrayInputStream(plainMessage);
				ByteArrayOutputStream out = new ByteArrayOutputStream();) {
			Cipher cipher = Cipher.getInstance(ALGORITHM_RSA);
			cipher.init(Cipher.ENCRYPT_MODE, publicKey);
			byte[] buf = new byte[ENCRYPT_BUFFER_SIZE];
			int readCount = 0;
			while ((readCount = in.read(buf)) != -1) {
				byte[] block = null;
				if (buf.length == readCount) {
					block = buf;
				} else {
					block = new byte[readCount];
					System.arraycopy(buf, 0, block, 0, readCount);
				}
				out.write(cipher.doFinal(block));
				out.flush();
			}
			return out.toByteArray();
		} catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException | IllegalBlockSizeException
				| BadPaddingException | IOException e) {
			throw new IllegalStateException(e);
		}
	}

	/**
	 * 生成公私秘钥对
	 * 
	 * @return
	 */
	public static KeyPair generateKeyPair() {
		try {
			KeyPairGenerator generator = KeyPairGenerator.getInstance(ALGORITHM_RSA);
			generator.initialize(KEY_SIZE);
			return generator.genKeyPair();
		} catch (NoSuchAlgorithmException e) {
			throw new IllegalStateException(e);
		}
	}

	public static String base64(String text) {
		return base64(text.getBytes(Charset.forName(DEFAULT_ENCODE)));
	}

	public static String base64(byte[] bytes) {
		return new String(Base64.encodeBase64(bytes), Charset.forName(DEFAULT_ENCODE));
	}

	public static byte[] decodeBase64(String base64Text) {
		return decodeBase64(base64Text.getBytes(Charset.forName(DEFAULT_ENCODE)));
	}

	public static byte[] decodeBase64(byte[] base64Bytes) {
		return Base64.decodeBase64(base64Bytes);
	}
}